# -*- coding: utf-8 -*-
import numpy as np
import cv2
import glob
import os

def undistortion(img, mtx, dist):
    h, w = img.shape[:2]
    # print('w is {}, h is {}'.format(w, h))
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))

    # print('newcameramatrix ', newcameramtx)
    newcameramtx = mtx
    dst = cv2.undistort(img, mtx, dist, None, newcameramtx)

    # crop the image
    x, y, w, h = roi
    print('w is {}, h is {}'.format(w, h))
    if roi != (0, 0, 0, 0):
        dst = dst[y:y + h, x:x + w]

    return dst

def calibrate():
    # 标注相机内参
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    Nx_cor = 9
    Ny_cor = 6

    objp = np.zeros((Nx_cor * Ny_cor, 3), np.float32)
    objp[:, :2] = np.mgrid[0:Nx_cor, 0:Ny_cor].T.reshape(-1, 2)# 此处认为棋盘格的大小为1cm
    #  objp[:, :2] = np.mgrid[0:Nx_cor * squareSize:squareSize, 0:Ny_cor * squareSize:squareSize].T.reshape(-1, 2)# 此处可以选定squareSize的棋盘大小
    objpoints = []  # 3d points in real world space
    imgpoints = []  # 2d points in image plane.
    images = glob.glob(r'文件名\\*.jpg')
    print(images)



    for fname in images:
        frame = cv2.imread(fname)
        # Our operations on the frame come here
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    
        ret, corners = cv2.findChessboardCorners(gray, (Nx_cor, Ny_cor), None)  # Find the corners
        print(ret)
        # If found, add object points, image points
        if ret == True:
            corners = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)
            objpoints.append(objp)
            imgpoints.append(corners)
            cv2.drawChessboardCorners(frame, (Nx_cor, Ny_cor), corners, ret)
            

       
    global mtx, dist

    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
    print(mtx, dist)

    mean_error = 0
    for i in range(len(objpoints)):
        imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
        error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        mean_error += error

    print ("total error: ", mean_error / len(objpoints))
        # # When everything done, release the capture

    np.savez('calibrate.npz', mtx=mtx, dist=dist[0:4])




if __name__ == '__main__':

    cap = cv2.VideoCapture(0)

    mtx = []
    dist = []


    try:
        npzfile = np.load('calibrate_test.npz')
        mtx = npzfile['mtx']
        dist = npzfile['dist']
    except IOError:
        calibrate()

    print('dist', dist[0:4])
    print(mtx)
#    images = glob.glob(r'PatternImage/*.jpg')
#    new_path = 'NewImage'
#    if not os.path.exists(new_path):
#        os.mkdir(new_path)
#    for fname in images:
#        print(fname)
#        frame = cv2.imread(fname)
#        frame = undistortion(frame, mtx, dist[0:4])
#        # Display the resulting frame
#        base_name = os.path.basename(fname)
#        cv2.imwrite(os.path.join(new_path, base_name), frame)
        


